
#include "mish_grad_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
const uint32_t BLOCK_SIZE = 32;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

    MishGradTilingData tiling;
    int32_t NUM = 12; //申请的内存块的个数（包括buffer和queue）
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    uint64_t ub_size; 
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
    auto aivNum = ascendcPlatform.GetCoreNumAiv();

    uint32_t totalDataNum = context->GetInputTensor(0)->GetShapeSize();
    // uint32_t inputNum = context->GetInputNum();

    // if(inputNum == 2)
    // {
    //     tiling.set_optional(2);
    // }
    // else if (inputNum == 3)
    // {
    //     tiling.set_optional(3);
    // }
    
    auto dt = context->GetInputDesc(0)->GetDataType();
    uint32_t sizeofdatatype;
    if (dt == ge::DT_FLOAT16) 
    {
        sizeofdatatype = 2;
    }
    else 
    {
        sizeofdatatype = 4;
    }

    uint32_t ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;

    uint32_t tiling_size = ((ub_size) / BLOCK_SIZE / 2 ) / NUM;

    tiling_size = tiling_size <= 8 ? tiling_size : tiling_size / 8 * 8; //256Byte对齐
    uint32_t block_size = tiling_size * ALIGN_NUM;
    // ub_data_num = 16;

    uint32_t totalLengthNumAligned = (totalDataNum + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;

    context->SetBlockDim(1);

    uint32_t tileNum = totalLengthNumAligned / block_size;

    uint32_t tileLength = 0;
    uint32_t lasttileLength = 0;
    // 
    if(totalLengthNumAligned % block_size == 0 || tileNum == 0) //reduce dim对齐之后，可以均分到各个tile，或不足一个tile
    {
        if(tileNum == 0)    //一次tiling的数据个数 大于 totalLengthAligned，也就是tile_num = 0时，tileLength=totalLengthAligned
        {
            tileNum = 1;
            tileLength = totalLengthNumAligned;
            lasttileLength = tileLength;
        }
        else //一次tiling的数据个数 小于 totalLengthAligned时，tileLength取最大值tiling_size * ALIGN_NUM
        {
            tileLength = block_size;
            lasttileLength = tileLength;
        }
        
    }
    else //totalLengthAligned对齐后不能均分到每次tiling
    {
        tileNum = tileNum + 1;
        tileLength = block_size; //tileLength 取最大值(ub)
        lasttileLength = totalLengthNumAligned % tileLength;
    }

    

    // printf("inputNum: %d\n", inputNum);
    // printf("ALIGN_NUM: %d\n", ALIGN_NUM);
    // printf("block_size: %d\n", block_size);
    // printf("totalDataNum: %d\n", totalDataNum);
    // printf("tileNum: %d\n", tileNum);
    // printf("tileLength: %d\n", tileLength);
    // printf("lasttileLength: %d\n", lasttileLength);


    tiling.set_ALIGN_NUM(ALIGN_NUM);
    tiling.set_totalLength(totalDataNum);
    tiling.set_tileNum(tileNum);
    tiling.set_tileLength(tileLength);
    tiling.set_lasttileLength(lasttileLength);
    

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());


  return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class MishGrad : public OpDef {
public:
    explicit MishGrad(const char* name) : OpDef(name)
    {
        this->Input("grad")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("tanhx")
            .ParamType(OPTIONAL)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(MishGrad);
}
